import logging
import abc
import copy
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from crowd_sim.envs.utils.action import ActionXY
import numpy as np
import itertools

class SocialTrainer(object):
    def __init__(self, value_estimator, state_predictor, memory, device, policy, writer, batch_size, optimizer_str, human_num,
                 reduce_sp_update_frequency, freeze_state_predictor, detach_state_predictor, share_graph_model):
        """
        Train the trainable model of a policy
        """
        self.value_estimator = value_estimator
        self.state_predictor = state_predictor
        self.device = device
        self.writer = writer
        self.target_policy = policy
        self.target_model = None
        self.criterion = nn.MSELoss().to(device)
        self.memory = memory
        self.data_loader = None
        self.batch_size = batch_size
        self.optimizer_str = optimizer_str
        self.reduce_sp_update_frequency = reduce_sp_update_frequency
        self.state_predictor_update_interval = human_num
        self.freeze_state_predictor = freeze_state_predictor
        self.detach_state_predictor = detach_state_predictor
        self.share_graph_model = share_graph_model
        self.v_optimizer_main = None
        self.v_optimizer_auxi = None
        self.s_optimizer = None

        # for value update
        self.gamma = 0.9
        self.time_step = 0.25
        self.v_pref = 1

    def update_target_model(self, target_model):
        self.target_model = copy.deepcopy(target_model)

    # def set_learning_rate(self, learning_rate):
    #     if self.optimizer_str == 'Adam':
    #         self.v_optimizer_main = optim.Adam(self.value_estimator.parameters(), lr=learning_rate)
    #         self.v_optimizer_auxi = optim.Adam(self.value_estimator.graph_model.tpcnns.parameters(), lr=learning_rate)
    #         if self.state_predictor.trainable:
    #             self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
    #     elif self.optimizer_str == 'SGD':
    #         self.v_optimizer_main = optim.SGD(self.value_estimator.parameters(), lr=learning_rate, momentum=0.9)
    #         self.v_optimizer_auxi = optim.SGD(self.value_estimator.graph_model.tpcnns.parameters(), lr=learning_rate, momentum=0.9)
    #         if self.state_predictor.trainable:
    #             self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate,momentum=0.9)
    #     else:
    #         raise NotImplementedError
    #
    #     if self.state_predictor.trainable:
    #         logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
    #             [name for name, param in list(self.value_estimator.named_parameters()) +
    #              list(self.state_predictor.named_parameters())]), self.optimizer_str))
    #     else:
    #         logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
    #             [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))

    def set_rl_learning_rate(self, learning_rate):   # 0.0005
        if self.optimizer_str == 'Adam':
            self.v_optimizer_main = optim.Adam(itertools.chain(self.value_estimator.parameters()), lr=learning_rate)
            self.v_optimizer_auxi = optim.Adam(itertools.chain(self.value_estimator.graph_model.tpcnns.parameters()), lr=learning_rate)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.v_optimizer_main = optim.SGD(itertools.chain(self.value_estimator.parameters()), lr=learning_rate, momentum=0.9)
            self.v_optimizer_auxi = optim.SGD(itertools.chain(self.value_estimator.graph_model.tpcnns.parameters(),self.value_estimator.graph_model.tpcnn_ouput.parameters()), lr=learning_rate, momentum=0.9)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate)
        else:
            raise NotImplementedError

        if self.state_predictor.trainable:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters()) +
                 list(self.state_predictor.named_parameters())]), self.optimizer_str))
        else:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))
    '''
    def optimize_epoch(self, num_epochs):
        if self.v_optimizer_main is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        for epoch in range(num_epochs):
            epoch_v_loss = 0
            epoch_s_loss = 0
            logging.debug('{}-th epoch starts'.format(epoch))

            update_counter = 0
            for data in self.data_loader:
                robot_states, human_states, actions, values, _, next_robot_state, next_human_states = data

                # optimize value estimator
                self.v_optimizer_main.zero_grad()
                actions = actions.to(self.device)
                outputs = self.value_estimator((robot_states, human_states)).gather(1, actions.unsqueeze(1))
                values = values.to(self.device)
                loss = self.criterion(outputs, values)
                loss.backward()
                self.v_optimizer.step()
                epoch_v_loss += loss.data.item()

                # optimize state predictor
                if self.state_predictor.trainable:
                    update_state_predictor = True
                    if update_counter % self.state_predictor_update_interval != 0:
                        update_state_predictor = False

                    if update_state_predictor:
                        self.s_optimizer.zero_grad()
                        _, next_human_states_est = self.state_predictor((robot_states, human_states), None)
                        loss = self.criterion(next_human_states_est, next_human_states)
                        loss.backward()
                        self.s_optimizer.step()
                        epoch_s_loss += loss.data.item()
                    update_counter += 1
                else:
                    _, next_human_states_est = self.state_predictor((robot_states, human_states), ActionXY(0, 0))
                    loss = self.criterion(next_human_states_est, next_human_states)
                    epoch_s_loss += loss.data.item()

            logging.debug('{}-th epoch ends'.format(epoch))
            self.writer.add_scalar('IL/epoch_v_loss', epoch_v_loss / len(self.memory), epoch)
            self.writer.add_scalar('IL/epoch_s_loss', epoch_s_loss / len(self.memory), epoch)
            logging.info('Average loss in epoch %d: %.2E, %.2E', epoch, epoch_v_loss / len(self.memory),
                         epoch_s_loss / len(self.memory))
        return
        '''

    def optimize_batch(self, num_batches, episode):
        if self.v_optimizer_main is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        RL_losses = 0
        auxiliary_losses = 0
        batch_count = 0
        average_inner_reward = 0
        self.target_model.value_network.eval()
        self.value_estimator.value_network.eval()
        for data in self.data_loader:
            batch_num = int(self.data_loader.sampler.num_samples // self.batch_size)
            robot_states, human_states, actions, _, rewards, next_robot_states, next_human_states, humans_label = data # robot_states [100,1,9]--->[100,1,8,9]

            # optimize value estimator


            actions = actions.to(self.device)
            # outputs = self.value_estimator((robot_states, human_states))
            # aaa = self.value_estimator((robot_states, human_states))
            # outputs = self.value_estimator((robot_states, human_states)).gather(1, actions.unsqueeze(1))
            robot_Q, predicted_human_trajectory = self.value_estimator((robot_states, human_states))
            outputs = robot_Q.gather(1, actions.unsqueeze(1))
            gamma_bar = pow(self.gamma, self.time_step * self.v_pref)
            max_next_Q_index = torch.max(self.value_estimator((next_robot_states, next_human_states))[0], dim=1)[1]
            next_Q_value = self.target_model((next_robot_states, next_human_states))[0].gather(1, max_next_Q_index.unsqueeze(1))

            # George also uses human trajectory prediction loss。
            # humans_label_position: [batch_size,8,human_num,5] 5是 [self.px, self.py, self.vx, self.vy, self.radius]
            # predicted_human_trajectory: [batch_size,5,8,human_num]这个5是为了表示xy轨迹的二元高斯分布，
            # predicted_human_trajectory = predicted_human_trajectory.permute(0, 2, 3, 1)  # [100,8,6,5]
            humans_label_position = humans_label[:, :, :, 0:2]  # [100,8,6,2]
            inner_reward = self.bivariate_loss(predicted_human_trajectory, humans_label_position)
            auxiliary_loss = inner_reward.mean()
            inner_reward = inner_reward.detach()

            for i in range(inner_reward.shape[0]):
                # auxiliary_loss.data[i] *= 0.02
                average_inner_reward += inner_reward.data[i].item()
            average_inner_reward /= inner_reward.shape[0]

            inner_reward_weight = abs((abs(rewards).mean().item() * 1) / abs(inner_reward).mean().item())
            rewards -= inner_reward_weight * inner_reward
            # for dqn
            # max_next_Q = torch.max(self.target_model((next_robot_states, next_human_states)), dim=1)[0]
            # next_Q_value= max_next_Q.unsqueeze(dim=1)
            target_values = rewards + next_Q_value * gamma_bar
            # target_values = rewards + gamma_bar * self.target_model((next_robot_states, next_human_states))

            # values = values.to(self.device)
            RL_loss = self.criterion(outputs, target_values)
            RL_losses += RL_loss.data.item()

            loss = RL_loss + auxiliary_loss * 0.0003

            # 1
            # loss.backward()
            # self.v_optimizer_main.step()
            # self.v_optimizer_auxi.step()

            #2
            self.v_optimizer_main.zero_grad()
            RL_loss.backward(retain_graph=True)
            self.v_optimizer_main.step()

            self.v_optimizer_auxi.zero_grad()
            auxiliary_loss.backward()
            self.v_optimizer_auxi.step()


            batch_count += 1
            if batch_count > num_batches or batch_count == batch_num:
                break

        average_rl_loss = RL_losses / num_batches
        average_auxiliary_losses = auxiliary_losses / num_batches
        # average_s_loss = s_losses / num_batches  # for state_predictor
        logging.info('Average loss : %.2E, %.2E', average_rl_loss, average_auxiliary_losses)
        average_inner_reward /= num_batches
        logging.info('inner reward: %.2E', average_inner_reward)
        self.writer.add_scalar('average_rl_loss', average_rl_loss, episode)
        self.writer.add_scalar('average_inner_reward', average_inner_reward, episode)
        self.writer.add_scalar('average_auxiliary_losses', average_auxiliary_losses, episode)
        self.value_estimator.value_network.train()
        return average_rl_loss, average_auxiliary_losses

    def bivariate_loss(self, V_pred, V_trgt):
        # V_pred: [predicted_sequence_length, human_num, 5]  这个5是mux, muy, sx, sy, corr分别代表什么？ 二元高斯分布里的概念
        # V_trgt: [predicted_sequence_length, human_num, 2]   这个2就是x和y的真值了
        # mux, muy, sx, sy, corr
        # assert V_pred.shape == V_trgt.shape
        normx = V_trgt[:, :, :, 0] - V_pred[:, :, :, 0]
        normy = V_trgt[:, :, :, 1] - V_pred[:, :, :, 1]

        sx = torch.exp(V_pred[:, :, :, 2])  # sx
        sy = torch.exp(V_pred[:, :, :, 3])  # sy
        corr = torch.tanh(V_pred[:, :, :, 4])  # corr

        sxsy = sx * sy

        z = (normx / sx) ** 2 + (normy / sy) ** 2 - 2 * ((corr * normx * normy) / sxsy)
        negRho = 1 - corr ** 2

        # Numerator
        result = torch.exp(-z / (2 * negRho))
        # Normalization factor
        denom = 2 * np.pi * (sxsy * torch.sqrt(negRho))

        # Final PDF calculation    PDF:概率密度函数(probability density function)
        result = result / denom

        # Numerical stability
        epsilon = 1e-20

        result = -torch.log(torch.clamp(result, min=epsilon))
        # result = torch.mean(result)
        result = torch.mean(result, dim=1, keepdim=True).mean(dim=2)

        return result


class TSRLTrainer(object):
    def __init__(self, value_estimator, state_predictor, memory, device, policy, writer, batch_size, optimizer_str, human_num,
                 reduce_sp_update_frequency, freeze_state_predictor, detach_state_predictor, share_graph_model):
        """
        Train the trainable model of a policy
        """
        self.value_estimator = value_estimator
        self.state_predictor = state_predictor
        self.device = device
        self.writer = writer
        self.target_policy = policy
        self.target_model = None
        self.criterion = nn.MSELoss().to(device)
        self.memory = memory
        self.data_loader = None
        self.batch_size = batch_size
        self.optimizer_str = optimizer_str
        self.reduce_sp_update_frequency = reduce_sp_update_frequency
        self.state_predictor_update_interval = human_num
        self.freeze_state_predictor = freeze_state_predictor
        self.detach_state_predictor = detach_state_predictor
        self.share_graph_model = share_graph_model
        self.v_optimizer = None
        self.s_optimizer = None

        # for value update
        self.gamma = 0.9
        self.time_step = 0.25
        self.v_pref = 1

    def update_target_model(self, target_model):
        self.target_model = copy.deepcopy(target_model)

    def set_learning_rate(self, learning_rate):
        if self.optimizer_str == 'Adam':
            self.v_optimizer = optim.Adam(self.value_estimator.parameters(), lr=learning_rate)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.v_optimizer = optim.SGD(self.value_estimator.parameters(), lr=learning_rate, momentum=0.9)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate,momentum=0.9)
        else:
            raise NotImplementedError

        if self.state_predictor.trainable:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters()) +
                 list(self.state_predictor.named_parameters())]), self.optimizer_str))
        else:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))

    def set_rl_learning_rate(self, learning_rate):
        if self.optimizer_str == 'Adam':
            self.v_optimizer = optim.Adam(self.value_estimator.parameters(), lr=learning_rate)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.v_optimizer = optim.SGD(self.value_estimator.parameters(), lr=learning_rate, momentum=0.9)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate)
        else:
            raise NotImplementedError

        if self.state_predictor.trainable:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters()) +
                 list(self.state_predictor.named_parameters())]), self.optimizer_str))
        else:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))

    def optimize_epoch(self, num_epochs):
        if self.v_optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        for epoch in range(num_epochs):
            epoch_v_loss = 0
            epoch_s_loss = 0
            logging.debug('{}-th epoch starts'.format(epoch))

            update_counter = 0
            for data in self.data_loader:
                robot_states, human_states, actions, values, _, next_robot_state, next_human_states = data

                # optimize value estimator
                self.v_optimizer.zero_grad()
                actions = actions.to(self.device)
                outputs = self.value_estimator((robot_states, human_states)).gather(1, actions.unsqueeze(1))
                values = values.to(self.device)
                loss = self.criterion(outputs, values)
                loss.backward()
                self.v_optimizer.step()
                epoch_v_loss += loss.data.item()

                # optimize state predictor
                if self.state_predictor.trainable:
                    update_state_predictor = True
                    if update_counter % self.state_predictor_update_interval != 0:
                        update_state_predictor = False

                    if update_state_predictor:
                        self.s_optimizer.zero_grad()
                        _, next_human_states_est = self.state_predictor((robot_states, human_states), None)
                        loss = self.criterion(next_human_states_est, next_human_states)
                        loss.backward()
                        self.s_optimizer.step()
                        epoch_s_loss += loss.data.item()
                    update_counter += 1
                else:
                    _, next_human_states_est = self.state_predictor((robot_states, human_states), ActionXY(0, 0))
                    loss = self.criterion(next_human_states_est, next_human_states)
                    epoch_s_loss += loss.data.item()

            logging.debug('{}-th epoch ends'.format(epoch))
            self.writer.add_scalar('IL/epoch_v_loss', epoch_v_loss / len(self.memory), epoch)
            self.writer.add_scalar('IL/epoch_s_loss', epoch_s_loss / len(self.memory), epoch)
            logging.info('Average loss in epoch %d: %.2E, %.2E', epoch, epoch_v_loss / len(self.memory),
                         epoch_s_loss / len(self.memory))
        return

    def optimize_batch(self, num_batches, episode):
        if self.v_optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        v_losses = 0
        s_losses = 0
        batch_count = 0
        self.target_model.value_network.eval()
        self.value_estimator.value_network.eval()
        for data in self.data_loader:
            batch_num = int(self.data_loader.sampler.num_samples // self.batch_size)
            robot_states, human_states, actions, _, rewards, next_robot_states, next_human_states = data

            # optimize value estimator
            self.v_optimizer.zero_grad()
            actions = actions.to(self.device)
            # outputs = self.value_estimator((robot_states, human_states))
            outputs = self.value_estimator((robot_states, human_states)).gather(1, actions.unsqueeze(1))
            gamma_bar = pow(self.gamma, self.time_step * self.v_pref)
            max_next_Q_index = torch.max(self.value_estimator((next_robot_states, next_human_states)), dim=1)[1]
            next_Q_value = self.target_model((next_robot_states, next_human_states)).gather(1, max_next_Q_index.unsqueeze(1))
            # for dqn
            # max_next_Q = torch.max(self.target_model((next_robot_states, next_human_states)), dim=1)[0]
            # next_Q_value= max_next_Q.unsqueeze(dim=1)
            target_values = rewards + next_Q_value * gamma_bar
            # target_values = rewards + gamma_bar * self.target_model((next_robot_states, next_human_states))

            # values = values.to(self.device)
            loss = self.criterion(outputs, target_values)
            loss.backward()
            self.v_optimizer.step()
            v_losses += loss.data.item()

            # optimize state predictor
            if self.state_predictor.trainable:
                update_state_predictor = True
                if self.freeze_state_predictor:
                    update_state_predictor = False
                elif self.reduce_sp_update_frequency and batch_count % self.state_predictor_update_interval == 0:
                    update_state_predictor = False

                if update_state_predictor:
                    self.s_optimizer.zero_grad()
                    _, next_human_states_est = self.state_predictor((robot_states, human_states), None,
                                                                    detach=self.detach_state_predictor)
                    loss = self.criterion(next_human_states_est, next_human_states)
                    loss.backward()
                    self.s_optimizer.step()
                    s_losses += loss.data.item()
            else:
                _, next_human_states_est = self.state_predictor((robot_states, human_states), None,
                                                                detach=self.detach_state_predictor)
                loss = self.criterion(next_human_states_est, next_human_states)
                s_losses += loss.data.item()


            batch_count += 1
            if batch_count > num_batches or batch_count == batch_num:
                break

        average_v_loss = v_losses / num_batches
        average_s_loss = s_losses / num_batches
        logging.info('Average loss : %.2E, %.2E', average_v_loss, average_s_loss)
        self.writer.add_scalar('RL/average_v_loss', average_v_loss, episode)
        self.writer.add_scalar('RL/average_s_loss', average_s_loss, episode)
        self.value_estimator.value_network.train()
        return average_v_loss, average_s_loss

class MPRLTrainer(object):
    def __init__(self, value_estimator, state_predictor, memory, device, policy, writer, batch_size, optimizer_str, human_num,
                 reduce_sp_update_frequency, freeze_state_predictor, detach_state_predictor, share_graph_model):
        """
        Train the trainable model of a policy
        """
        self.value_estimator = value_estimator
        self.state_predictor = state_predictor
        self.device = device
        self.writer = writer
        self.target_policy = policy
        self.target_model = None
        self.criterion = nn.MSELoss().to(device)
        self.memory = memory
        self.data_loader = None
        self.batch_size = batch_size
        self.optimizer_str = optimizer_str
        self.reduce_sp_update_frequency = reduce_sp_update_frequency
        self.state_predictor_update_interval = human_num
        self.freeze_state_predictor = freeze_state_predictor
        self.detach_state_predictor = detach_state_predictor
        self.share_graph_model = share_graph_model
        self.v_optimizer = None
        self.s_optimizer = None

        # for value update
        self.gamma = 0.9
        self.time_step = 0.25
        self.v_pref = 1

    def update_target_model(self, target_model):
        self.target_model = copy.deepcopy(target_model)

    def set_learning_rate(self, learning_rate):
        if self.optimizer_str == 'Adam':
            # self.v_optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.value_estimator.parameters()), lr=learning_rate)
            self.v_optimizer = optim.Adam(self.value_estimator.parameters(), lr=learning_rate)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.v_optimizer = optim.SGD(self.value_estimator.parameters(), lr=learning_rate, momentum=0.9)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate, momentum=0.9)
        else:
            raise NotImplementedError

        if self.state_predictor.trainable:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters()) +
                 list(self.state_predictor.named_parameters())]), self.optimizer_str))
        else:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))

    def set_rl_learning_rate(self, learning_rate):
        if self.optimizer_str == 'Adam':
            self.v_optimizer = optim.Adam(self.value_estimator.parameters(), lr=learning_rate)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.Adam(self.state_predictor.parameters(), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.v_optimizer = optim.SGD(self.value_estimator.parameters(), lr=learning_rate, momentum=0.9)
            if self.state_predictor.trainable:
                self.s_optimizer = optim.SGD(self.state_predictor.parameters(), lr=learning_rate)
        else:
            raise NotImplementedError

        if self.state_predictor.trainable:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters()) +
                 list(self.state_predictor.named_parameters())]), self.optimizer_str))
        else:
            logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
                [name for name, param in list(self.value_estimator.named_parameters())]), self.optimizer_str))

    def optimize_epoch(self, num_epochs):
        if self.v_optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        for epoch in range(num_epochs):
            epoch_v_loss = 0
            epoch_s_loss = 0
            logging.debug('{}-th epoch starts'.format(epoch))

            update_counter = 0
            for data in self.data_loader:
                robot_states, human_states, actions, values, _, next_robot_state, next_human_states = data

                # optimize value estimator
                self.v_optimizer.zero_grad()
                actions = actions.to(self.device)
                outputs = self.value_estimator((robot_states, human_states))
                values = values.to(self.device)
                loss = self.criterion(outputs, values)
                loss.backward()
                self.v_optimizer.step()
                epoch_v_loss += loss.data.item()

                # optimize state predictor
                if self.state_predictor.trainable:
                    update_state_predictor = True
                    if update_counter % self.state_predictor_update_interval != 0:
                        update_state_predictor = False

                    if update_state_predictor:
                        self.s_optimizer.zero_grad()
                        _, next_human_states_est = self.state_predictor((robot_states, human_states), None)
                        loss = self.criterion(next_human_states_est, next_human_states)
                        loss.backward()
                        self.s_optimizer.step()
                        epoch_s_loss += loss.data.item()
                    update_counter += 1
                else:
                    _, next_human_states_est = self.state_predictor((robot_states, human_states), ActionXY(0, 0))
                    loss = self.criterion(next_human_states_est, next_human_states)
                    epoch_s_loss += loss.data.item()

            logging.debug('{}-th epoch ends'.format(epoch))
            self.writer.add_scalar('IL/epoch_v_loss', epoch_v_loss / len(self.memory), epoch)
            self.writer.add_scalar('IL/epoch_s_loss', epoch_s_loss / len(self.memory), epoch)
            logging.info('Average loss in epoch %d: %.2E, %.2E', epoch, epoch_v_loss / len(self.memory),
                         epoch_s_loss / len(self.memory))
        return

    def optimize_batch(self, num_batches, episode):
        if self.v_optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True)
        v_losses = 0
        s_losses = 0
        batch_count = 0
        self.target_model.value_network.eval()
        self.value_estimator.value_network.eval()
        for data in self.data_loader:
            batch_num = int(self.data_loader.sampler.num_samples // self.batch_size)
            robot_states, human_states, actions, _, rewards, next_robot_states, next_human_states = data

            # optimize value estimator
            self.v_optimizer.zero_grad()
            actions = actions.to(self.device)
            # outputs = self.value_estimator((robot_states, human_states))
            outputs = self.value_estimator((robot_states, human_states))
            gamma_bar = pow(self.gamma, self.time_step * self.v_pref)
            next_value = self.target_model((next_robot_states, next_human_states))
            target_values = rewards + next_value * gamma_bar
            # values = values.to(self.device)
            loss = self.criterion(outputs, target_values)
            loss.backward()
            self.v_optimizer.step()
            v_losses += loss.data.item()

            # optimize state predictor
            if self.state_predictor.trainable:
                update_state_predictor = True
                if self.freeze_state_predictor:
                    update_state_predictor = False
                elif self.reduce_sp_update_frequency and batch_count % self.state_predictor_update_interval == 0:
                    update_state_predictor = False

                if update_state_predictor:
                    self.s_optimizer.zero_grad()
                    _, next_human_states_est = self.state_predictor((robot_states, human_states), None,
                                                                    detach=self.detach_state_predictor)
                    loss = self.criterion(next_human_states_est, next_human_states)
                    loss.backward()
                    self.s_optimizer.step()
                    s_losses += loss.data.item()
            else:
                _, next_human_states_est = self.state_predictor((robot_states, human_states), None,
                                                                detach=self.detach_state_predictor)
                loss = self.criterion(next_human_states_est, next_human_states)
                s_losses += loss.data.item()
            batch_count += 1
            if batch_count > num_batches or batch_count == batch_num:
                break

        average_v_loss = v_losses / num_batches
        average_s_loss = s_losses / num_batches
        logging.info('Average loss : %.2E, %.2E', average_v_loss, average_s_loss)
        self.writer.add_scalar('RL/average_v_loss', average_v_loss, episode)
        self.writer.add_scalar('RL/average_s_loss', average_s_loss, episode)
        self.value_estimator.value_network.train()
        return average_v_loss, average_s_loss

class VNRLTrainer(object):
    def __init__(self, model, memory, device, policy, batch_size, optimizer_str, writer):
        """
        Train the trainable model of a policy
        """
        self.model = model
        self.device = device
        self.policy = policy
        self.target_model = None
        self.criterion = nn.MSELoss().to(device)
        self.memory = memory
        self.data_loader = None
        self.batch_size = batch_size
        self.optimizer_str = optimizer_str
        self.optimizer = None
        self.writer = writer

        # for value update
        self.gamma = 0.9
        self.time_step = 0.25
        self.v_pref = 1

    def update_target_model(self, target_model):
        self.target_model = copy.deepcopy(target_model)

    def set_rl_learning_rate(self, learning_rate):
        if self.optimizer_str == 'Adam':
            self.optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=learning_rate)
        elif self.optimizer_str == 'SGD':
            self.optimizer = optim.SGD(self.model.parameters(), lr=learning_rate, momentum=0.9)
        else:
            raise NotImplementedError
        logging.info('Lr: {} for parameters {} with {} optimizer'.format(learning_rate, ' '.join(
            [name for name, param in self.model.named_parameters()]), self.optimizer_str))

    def optimize_epoch(self, num_epochs):
        if self.optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True, collate_fn=pad_batch)
        average_epoch_loss = 0
        for epoch in range(num_epochs):
            epoch_loss = 0
            logging.debug('{}-th epoch starts'.format(epoch))
            for data in self.data_loader:
                inputs, values, _, _ = data
                self.optimizer.zero_grad()
                outputs = self.model(inputs)
                values = values.to(self.device)
                loss = self.criterion(outputs, values)
                loss.backward()
                self.optimizer.step()
                epoch_loss += loss.data.item()
            logging.debug('{}-th epoch ends'.format(epoch))
            average_epoch_loss = epoch_loss / len(self.memory)
            self.writer.add_scalar('IL/average_epoch_loss', average_epoch_loss, epoch)
            logging.info('Average loss in epoch %d: %.2E', epoch, average_epoch_loss)

        return average_epoch_loss

    def optimize_batch(self, num_batches, episode=None):
        if self.optimizer is None:
            raise ValueError('Learning rate is not set!')
        if self.data_loader is None:
            self.data_loader = DataLoader(self.memory, self.batch_size, shuffle=True, collate_fn=pad_batch)
        losses = 0
        batch_count = 0
        for data in self.data_loader:
            inputs, _, rewards, next_states = data
            self.optimizer.zero_grad()
            outputs = self.model(inputs)

            gamma_bar = pow(self.gamma, self.time_step * self.v_pref)
            target_values = rewards + gamma_bar * self.target_model(next_states)

            loss = self.criterion(outputs, target_values)
            loss.backward()
            self.optimizer.step()
            losses += loss.data.item()
            batch_count += 1
            if batch_count > num_batches:
                break

        average_loss = losses / num_batches
        logging.info('Average loss : %.2E', average_loss)

        return average_loss


def pad_batch(batch):
    """
    args:
        batch - list of (tensor, label)
    return:
        xs - a tensor of all examples in 'batch' after padding
        ys - a LongTensor of all labels in batch
    """
    def sort_states(position):
        # sort the sequences in the decreasing order of length
        sequences = sorted([x[position] for x in batch], reverse=True, key=lambda t: t.size()[0])
        packed_sequences = torch.nn.utils.rnn.pack_sequence(sequences)
        return torch.nn.utils.rnn.pad_packed_sequence(packed_sequences, batch_first=True)

    states = sort_states(0)
    values = torch.cat([x[1] for x in batch]).unsqueeze(1)
    rewards = torch.cat([x[2] for x in batch]).unsqueeze(1)
    next_states = sort_states(3)

    return states, values, rewards, next_states
